import logging

import ray

from .remote_task import RayRemoteRunnableTask
from calc_manager.base.runnable import Runnable
from calc_manager.base.executor import TaskExecutor
from typing import List, Optional, Dict, Any, cast
from calc_manager.base.task import Task
from ray.actor import ActorHandle  # 导入 Ray 的 ActorHandle 类型


logger = logging.getLogger(__name__)


class RayManager(TaskExecutor):
    def __init__(self, address=None, ray_init_kwargs={}):
        self.address = address
        self.init_kwargs = ray_init_kwargs
        logger.info(f"RayManager initialized with address={address}, ray_init_kwargs={ray_init_kwargs}")
        self.initialize()

    def initialize(self):
        if not ray.is_initialized():
            ray.init(address=self.address, **self.init_kwargs )
            logger.info(f"Connected to Ray cluster at {self.address}")
        else:
            pass

    def shutdown(self):
        if ray.is_initialized():
            ray.shutdown()
            logger.info("Ray cluster disconnecteggd.")

    def invoke(self, ray_remote_task: Task, ctx: Optional[Dict[str, Any]] = None):
        self.initialize()
        
        # 将 ray_remote_task 转换为 ActorHandle 类型
        ray_actor = cast(ActorHandle, ray_remote_task)
        
        try:
            if ctx:
                ray.get(ray_actor.set_ctx.remote(ctx))
                logger.debug(f"task {ray_remote_task} set_ctx finished")

            result_remote = ray_actor.execute.remote()
            result = ray.get(result_remote)
            logger.debug(f"task {ray_remote_task} get_result finished")
            ctx = ray.get(ray_actor.get_ctx.remote())
            logger.debug(f"task {ray_remote_task} get_ctx finished")
            return (result, ctx)
        except Exception as err:
            #logger.error(f"Failed to run task {ray_remote_task}: the error is :\n {err}")  
            raise RuntimeError(f"Failed to run task due to: {err}") from err
        finally:
            ray_actor.dispose.remote() #close actor
            pass

    def invoke_tasks(self, tasks: List[RayRemoteRunnableTask]):
        results: Dict[str, Any] = {}
        current_ctx: Dict[str, Any] = {}
        for task in tasks:
            try:
                ret, ctx = self.invoke(task, current_ctx)
            except Exception as err:
                #logger.error("tasks invoke failed, err: %s", err)
                for t in tasks:
                    ray_actor = cast(ActorHandle, t)
                    ray_actor.dispose.remote()
                raise err

            current_ctx = ctx

            # If ret is a dict, merge it into results
            if isinstance(ret, dict):
                results.update(ret)
            else:
                # Use task name as key if ret is not a dict
                results[task.task_info.task_name] = ret

        return (results, current_ctx)
